import asyncio
import logging
import os
import pyodbc
from pathlib import Path
from mcp.server import Server
from mcp.types import Resource, Tool, TextContent
from pydantic import AnyUrl
from mcp.server.stdio import stdio_server
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("access_mcp_server")

class AccessDatabaseManager:
    """管理Access数据库连接和操作的类"""
    
    def __init__(self, db_path: str):
        """
        初始化Access数据库管理器
        
        Args:
            db_path: Access数据库文件路径
        """
        self.db_path = Path(db_path)
        self.connection_string = self._build_connection_string()
        
    def _build_connection_string(self) -> str:
        """构建Access数据库连接字符串"""
        # 使用Microsoft Access Driver (32-bit)或(64-bit)
        # 根据系统架构选择合适的驱动
        drivers = [
            "Microsoft Access Driver (*.mdb, *.accdb)",
            "{Microsoft Access Driver (*.mdb, *.accdb)}",
            "Microsoft Access Driver (*.mdb)",
            "{Microsoft Access Driver (*.mdb)}"
        ]
        
        # 尝试找到可用的驱动
        available_drivers = [driver for driver in pyodbc.drivers() if "Access" in driver]
        logger.info(f"Available Access drivers: {available_drivers}")
        
        if not available_drivers:
            raise RuntimeError("No Microsoft Access driver found. Please install Microsoft Access Database Engine.")
        
        # 使用第一个可用的Access驱动
        driver = available_drivers[0]
        
        # 构建连接字符串
        conn_str = f"Driver={{{driver}}};DBQ={self.db_path.absolute()};"
        logger.info(f"Using connection string: {conn_str}")
        return conn_str
    
    def get_connection(self):
        """获取数据库连接"""
        try:
            conn = pyodbc.connect(self.connection_string)
            logger.info(f"Successfully connected to Access database: {self.db_path}")
            return conn
        except Exception as e:
            logger.error(f"Failed to connect to Access database: {e}")
            raise
    
    def get_tables(self) -> list[str]:
        """获取数据库中的所有表名"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 尝试多种方法获取表名
            tables = []
            
            # 方法1: 使用系统表查询（需要权限）
            try:
                cursor.execute("SELECT Name FROM MSysObjects WHERE Type=1 AND Flags=0")
                tables = [row[0] for row in cursor.fetchall()]
                logger.info(f"Method 1 - Found tables: {tables}")
            except Exception as e:
                logger.warning(f"Method 1 failed: {e}")
            
            # 方法2: 使用系统表查询（简化版本）
            if not tables:
                try:
                    cursor.execute("SELECT Name FROM MSysObjects WHERE Type=1")
                    tables = [row[0] for row in cursor.fetchall()]
                    logger.info(f"Method 2 - Found tables: {tables}")
                except Exception as e:
                    logger.warning(f"Method 2 failed: {e}")
            
            # 方法3: 使用INFORMATION_SCHEMA（如果支持）
            if not tables:
                try:
                    cursor.execute("SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE='BASE TABLE'")
                    tables = [row[0] for row in cursor.fetchall()]
                    logger.info(f"Method 3 - Found tables: {tables}")
                except Exception as e:
                    logger.warning(f"Method 3 failed: {e}")
            
            # 方法4: 尝试直接查询已知的表名（通过错误信息推断）
            if not tables:
                try:
                    # 尝试一些常见的表名
                    common_tables = ['Customers', 'Orders', 'Products', 'Employees', 'Suppliers', 'Categories']
                    for table_name in common_tables:
                        try:
                            cursor.execute(f"SELECT TOP 1 * FROM [{table_name}]")
                            tables.append(table_name)
                        except:
                            pass
                    logger.info(f"Method 4 - Found tables: {tables}")
                except Exception as e:
                    logger.warning(f"Method 4 failed: {e}")
            
            cursor.close()
            conn.close()
            
            logger.info(f"Final tables found: {tables}")
            return tables
            
        except Exception as e:
            logger.error(f"Error getting tables: {e}")
            return []
    
    def get_table_schema(self, table_name: str) -> list[tuple]:
        """获取表结构信息"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 获取表结构
            cursor.execute(f"SELECT * FROM [{table_name}] WHERE 1=0")
            columns = [(desc[0], desc[1].__name__) for desc in cursor.description]
            
            cursor.close()
            conn.close()
            
            return columns
            
        except Exception as e:
            logger.error(f"Error getting schema for table {table_name}: {e}")
            return []
    
    def execute_query(self, query: str, limit: int = 100) -> tuple[list, list]:
        """执行SQL查询"""
        try:
            conn = self.get_connection()
            cursor = conn.cursor()
            
            # 添加LIMIT子句（Access使用TOP）
            if "SELECT" in query.upper() and "TOP" not in query.upper():
                # 简单的TOP添加逻辑（实际应用中可能需要更复杂的解析）
                if "ORDER BY" in query.upper():
                    # 如果有ORDER BY，在ORDER BY前添加TOP
                    parts = query.upper().split("ORDER BY")
                    query = f"{parts[0]} TOP {limit} * FROM {query.split('FROM')[1].split('ORDER BY')[0]} ORDER BY{parts[1]}"
                else:
                    # 如果没有ORDER BY，在FROM后添加TOP
                    parts = query.upper().split("FROM")
                    query = f"{parts[0]} TOP {limit} * FROM{parts[1]}"
            
            cursor.execute(query)
            
            # 获取列名
            columns = [desc[0] for desc in cursor.description] if cursor.description else []
            
            # 获取数据
            rows = cursor.fetchall()
            
            cursor.close()
            conn.close()
            
            return columns, rows
            
        except Exception as e:
            logger.error(f"Error executing query: {e}")
            raise

# 初始化Access数据库管理器
def get_access_manager() -> AccessDatabaseManager:
    """获取Access数据库管理器实例"""
    # 从环境变量或默认路径获取数据库文件路径
    db_path = os.getenv("ACCESS_DB_PATH", "ERP2.accdb")
    
    # 如果路径是相对路径，转换为绝对路径
    if not Path(db_path).is_absolute():
        db_path = Path.cwd() / db_path
    
    logger.info(f"Using Access database: {db_path}")
    
    if not Path(db_path).exists():
        raise FileNotFoundError(f"Access database file not found: {db_path}")
    
    return AccessDatabaseManager(str(db_path))

# 初始化服务器
app = Server("access_mcp_server")

@app.list_resources()
async def list_resources() -> list[Resource]:
    """列出Access数据库中的表作为资源"""
    try:
        access_mgr = get_access_manager()
        tables = access_mgr.get_tables()
        
        resources = []
        for table in tables:
            # 获取表结构信息
            schema = access_mgr.get_table_schema(table)
            schema_info = ", ".join([f"{col[0]}({col[1]})" for col in schema])
            
            resources.append(
                Resource(
                    uri=AnyUrl(f"access://{table}/data"),
                    name=f"Table: {table}",
                    mimeType="text/plain",
                    description=f"Data in table: {table}. Schema: {schema_info}"
                )
            )
        
        logger.info(f"Returning {len(resources)} resources")
        return resources
        
    except Exception as e:
        logger.error(f"Failed to list resources: {str(e)}")
        return []

@app.read_resource()
async def read_resource(uri: AnyUrl) -> str:
    """读取表内容"""
    uri_str = str(uri)
    logger.info(f"Reading resource: {uri_str}")
    
    if not uri_str.startswith("access://"):
        raise ValueError(f"Invalid URI scheme: {uri_str}")
        
    parts = uri_str[9:].split('/')  # 去掉 "access://" 前缀
    table = parts[0]
    
    try:
        access_mgr = get_access_manager()
        columns, rows = access_mgr.execute_query(f"SELECT * FROM [{table}]", limit=50)
        
        if not columns:
            return "No data found or table is empty"
        
        # 格式化输出
        result_lines = []
        
        # 添加列标题
        header = "\t".join(str(col) for col in columns)
        result_lines.append(header)
        result_lines.append("-" * len(header))
        
        # 添加数据行
        for row in rows:
            row_values = []
            for val in row:
                if val is None:
                    row_values.append("NULL")
                elif isinstance(val, bytes):
                    row_values.append(val.decode('utf-8', errors='replace'))
                elif isinstance(val, bool):
                    row_values.append(str(int(val)))
                elif isinstance(val, (int, float)):
                    row_values.append(str(val))
                else:
                    try:
                        row_values.append(str(val).strip())
                    except Exception as e:
                        logger.error(f"Error converting value {val}: {e}")
                        row_values.append("ERROR")
            result_lines.append("\t".join(row_values))
        
        return "\n".join(result_lines)
                
    except Exception as e:
        logger.error(f"Database error reading resource {uri}: {str(e)}")
        raise RuntimeError(f"Database error: {str(e)}")

@app.list_tools()
async def list_tools() -> list[Tool]:
    """列出可用的Access数据库工具"""
    logger.info("Listing tools...")
    return [
        Tool(
            name="execute_sql",
            description="Execute an SQL query on the Access database",
            inputSchema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The SQL query to execute"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "Maximum number of rows to return (default: 100)",
                        "default": 100
                    }
                },
                "required": ["query"]
            }
        ),
        Tool(
            name="list_tables",
            description="List all tables in the Access database",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        ),
        Tool(
            name="get_table_info",
            description="Get detailed information about a specific table",
            inputSchema={
                "type": "object",
                "properties": {
                    "table_name": {
                        "type": "string",
                        "description": "Name of the table to get information about"
                    }
                },
                "required": ["table_name"]
            }
        )
    ]

@app.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    """执行Access数据库工具"""
    try:
        access_mgr = get_access_manager()
        
        if name == "execute_sql":
            query = arguments.get("query")
            limit = arguments.get("limit", 100)
            
            if not query:
                raise ValueError("Query is required")
            
            columns, rows = access_mgr.execute_query(query, limit)
            
            if not columns:
                return [TextContent(type="text", text="Query executed successfully. No results to display.")]
            
            # 格式化输出
            output_lines = []
            
            # 添加标题
            header = "\t".join(str(col) for col in columns)
            output_lines.append(header)
            output_lines.append("-" * len(header))
            
            # 添加数据行
            for row in rows:
                row_values = []
                for val in row:
                    if val is None:
                        row_values.append("NULL")
                    elif isinstance(val, bytes):
                        row_values.append(val.decode('utf-8', errors='replace'))
                    elif isinstance(val, bool):
                        row_values.append(str(int(val)))
                    elif isinstance(val, (int, float)):
                        row_values.append(str(val))
                    else:
                        try:
                            row_values.append(str(val).strip())
                        except Exception as e:
                            logger.error(f"Error converting value {val}: {e}")
                            row_values.append("ERROR")
                output_lines.append("\t".join(row_values))
            
            result = "\n".join(output_lines)
            return [TextContent(type="text", text=result)]
            
        elif name == "list_tables":
            tables = access_mgr.get_tables()
            if not tables:
                return [TextContent(type="text", text="No tables found in the database.")]
            
            result = "Tables in the database:\n" + "\n".join([f"- {table}" for table in tables])
            return [TextContent(type="text", text=result)]
            
        elif name == "get_table_info":
            table_name = arguments.get("table_name")
            if not table_name:
                raise ValueError("Table name is required")
            
            schema = access_mgr.get_table_schema(table_name)
            if not schema:
                return [TextContent(type="text", text=f"Table '{table_name}' not found or no schema information available.")]
            
            # 获取表的行数
            try:
                columns, rows = access_mgr.execute_query(f"SELECT COUNT(*) as row_count FROM [{table_name}]")
                row_count = rows[0][0] if rows else 0
            except:
                row_count = "Unknown"
            
            result_lines = [f"Table: {table_name}"]
            result_lines.append(f"Row count: {row_count}")
            result_lines.append("\nColumns:")
            result_lines.append("Name\t\tType")
            result_lines.append("-" * 30)
            
            for col_name, col_type in schema:
                result_lines.append(f"{col_name}\t\t{col_type}")
            
            return [TextContent(type="text", text="\n".join(result_lines))]
            
        else:
            raise ValueError(f"Unknown tool: {name}")
                
    except Exception as e:
        logger.error(f"Error executing tool {name}: {e}", exc_info=True)
        return [TextContent(type="text", text=f"Error executing {name}: {str(e)}")]

async def main():
    """主入口点，运行MCP服务器"""
    logger.info("Starting Access MCP server...")
    
    try:
        # 测试数据库连接
        access_mgr = get_access_manager()
        tables = access_mgr.get_tables()
        logger.info(f"Successfully connected to Access database. Found {len(tables)} tables.")
        
        async with stdio_server() as (read_stream, write_stream):
            await app.run(
                read_stream,
                write_stream,
                app.create_initialization_options()
            )
    except Exception as e:
        logger.error(f"Server error: {str(e)}", exc_info=True)
        raise

if __name__ == "__main__":
    asyncio.run(main())
